package com.eobard.obfuscate.core;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import com.eobard.adapter.BaseVisitorAdapter;
import com.eobard.config.obfuscate.ObfuscatorConfig;
import com.eobard.enums.RunningEnum;
import com.eobard.obfuscate.IObfuscator;
import com.eobard.params.IMethod;
import com.eobard.params.OverrideMethodArgumentParams;
import com.eobard.utils.CollectionUtils;
import com.eobard.utils.ExceptionUtils;
import com.eobard.utils.ObfuscatorUtils;
import com.eobard.utils.RandomExpression;
import com.github.javaparser.StaticJavaParser;
import com.github.javaparser.ast.CompilationUnit;
import com.github.javaparser.ast.Modifier;
import com.github.javaparser.ast.NodeList;
import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
import com.github.javaparser.ast.body.FieldDeclaration;
import com.github.javaparser.ast.body.MethodDeclaration;
import com.github.javaparser.ast.body.Parameter;
import com.github.javaparser.ast.expr.Expression;
import com.github.javaparser.ast.expr.MarkerAnnotationExpr;
import com.github.javaparser.ast.expr.MethodCallExpr;
import com.github.javaparser.ast.stmt.Statement;
import com.github.javaparser.ast.type.Type;
import com.github.javaparser.resolution.declarations.ResolvedMethodDeclaration;
import com.github.javaparser.resolution.types.ResolvedReferenceType;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @Author: divenswu(核心代码)、Eobard Gu(重构代码)
 * @Date: 2025-04-27
 * @Description: 修改接口与实现类方法参数的核心混淆类
 */
@Component
public class OverrideMethodArgumentObfuscator extends BaseVisitorAdapter<OverrideMethodArgumentParams> implements IObfuscator {
    @Value("${obfuscate.enable.typeSolver.logging:false}")
    private boolean print;
    /**
     * 随机参数
     **/
    private final Random random = new Random();
    /**
     * 源码构建树
     **/
    private String sourceRoot;
    private final Map<String, CompilationUnit> fileMap = new HashMap<>();
    private static List<MethodDeclaration> tempStaticMethods = new ArrayList<>();
    /**
     * 保存当前重构方法的分组信息
     **/
    private final List<Set<String>> ancestorMethodList = new ArrayList<>();
    private final Map<String, List<MethodCallExpr>> methodCallExprMap = new ConcurrentHashMap<>();
    private final Map<String, MethodDeclaration> methodDeclarationMap = new ConcurrentHashMap<>();


    //数据初始化代码
    @Override
    public void initialize(ObfuscatorConfig config) {
        //清空上一次的历史记录
        CollectionUtils.clear(fileMap, methodCallExprMap, methodDeclarationMap);
        CollectionUtils.clear(tempStaticMethods, ancestorMethodList);


        ExceptionUtils.checkNonNull(
                List.of(config.getSourceRoot(), config.getOverrideMethodParams().getBasePackage()),
                List.of("请先设置项目路径!", "请先配置【接口与实现类方法参数混淆器】的basePackage!")
        );

        this.sourceRoot = config.getSourceRoot();
        initSolver(sourceRoot, config.getSourceExt());

        //获取源码文件
        ObfuscatorUtils.loopFiles(fileMap, sourceRoot);
        //获取模版里的静态函数
        tempStaticMethods = ObfuscatorUtils.getStaticTemplateMethods(staticMethodsFilePath);

    }

    //核心混淆代码
    @Override
    public void obfuscate(ObfuscatorConfig config) {
        OverrideMethodArgumentParams methodParams = config.getOverrideMethodParams();
        //计算当前需要修改的函数方法
        this.fileMap.forEach((path, cu) -> cu.accept(this, methodParams));
        //添加自定义的静态函数，制作扰乱的工具函数：
        List<IMethod.StaticMethod> newStaticMethods = randomStaticMethods(methodParams.getBasePackage());
        //进行代码的混淆逻辑
        for (Set<String> ancestorMethods : ancestorMethodList) {
            //1.随机获取一个函数, 对函数添加参数
            IMethod.StaticMethod staticMethod = newStaticMethods.get(random.nextInt(newStaticMethods.size()));
            //2.给当前的类添加一个loggerPrint函数
            for (String ancestorMethod : ancestorMethods) {
                //2.1.给当前的类添加一个loggerPrint函数
                MethodDeclaration methodDeclaration = methodDeclarationMap.get(ancestorMethod);
                methodDeclaration.findCompilationUnit().ifPresent(tcu -> {
                    tcu.findFirst(ClassOrInterfaceDeclaration.class).ifPresent(cif -> {
                        if (!cif.isInterface()) {
                            List<MethodDeclaration> methods = cif.getMethodsByName(RandomExpression.getLoggerPrintName());
                            if (methods.isEmpty()) {
                                MethodDeclaration m = cif.addMethod(RandomExpression.getLoggerPrintName(), Modifier.Keyword.PUBLIC, Modifier.Keyword.STATIC);
                                m.addParameter(new Parameter(StaticJavaParser.parseType("Object"), "values").setVarArgs(true));
//                                String body = "{try {logger.finest(\"" + RandomExpression.getName() + "\"+Arrays.toString(values));}catch (Exception ignored){}}";
                                String body = IMethod.getLoggerPrintBody(config);
                                m.setBody(StaticJavaParser.parseBlock(body));
                            }
                        }
                    });
                });
                //2.2.添加上面随机获取函数导入类
                methodDeclaration.findCompilationUnit().ifPresent(tcu -> {
                    tcu.findFirst(ClassOrInterfaceDeclaration.class).ifPresent(cif -> {
                        //如果是普通类
                        if (!cif.isInterface()) {
//                            tcu.addImport("java.util.logging.Logger");
                            tcu.addImport(IMethod.getLoggerImport(config));
                            tcu.addImport("java.util.Arrays");
                            tcu.addImport(staticMethod.packageName + "." + staticMethod.className);

                            if (config.getRunningStatus() == RunningEnum.JAVA_PROJECT) {
                                //检查是否含有logger字段
                                boolean hasLogger = cif.getFields()
                                        .stream()
                                        .flatMap(fd -> fd.getVariables().stream())
                                        .anyMatch(var -> var.getNameAsString().equals("logger"));

                                if (!hasLogger) {
                                    String className = cif.getNameAsString();
                                    FieldDeclaration fieldDeclaration = StaticJavaParser.parseBodyDeclaration(
                                            "private static final Logger logger=Logger.getLogger(" + className + ".class.getName());"
                                    ).asFieldDeclaration();
                                    cif.addMember(fieldDeclaration);
                                }
                            }
                        }
                    });
                });
            }
            //3.获取重构函数需要插入参数的位置
            NodeList<Parameter> staticParams = staticMethod.md.getParameters();
            String ancestorMethodFirst = new ArrayList<>(ancestorMethods).get(0);
            MethodDeclaration methodDeclarationFirst = methodDeclarationMap.get(ancestorMethodFirst);
            NodeList<Parameter> parametersFirst = methodDeclarationFirst.getParameters();
            boolean lastParamIsVarArg = !parametersFirst.isEmpty() && parametersFirst.get(parametersFirst.size() - 1).isVarArgs();
            int paramNum = Math.max(random.nextInt(methodParams.getMaxAddParamNum()), methodParams.getMinAddParamNum());
            List<String> tempList = parametersFirst.stream().map(Parameter::getNameAsString).toList();
            List<Type> paramsTypeList = new ArrayList<>();
            List<Integer> insertIndexList = new ArrayList<>();
            for (int i = 0; i < Math.max(paramNum, staticParams.size()); i++) {
                insertIndexList.add(random.nextInt(tempList.size() + (lastParamIsVarArg ? 0 : 1)));
                paramsTypeList.add(i < staticParams.size() ? staticParams.get(i).getType() : RandomExpression.randomType());
            }
            //4.对所有的重构函数添加参数
            for (String ancestorMethod : ancestorMethods) {
                //4.1.对所有的重构函数添加参数
                List<IMethod.AddTypeInfo> params = new ArrayList<>();
                MethodDeclaration methodDeclaration = methodDeclarationMap.get(ancestorMethod);
                boolean hasOverride = methodDeclaration.getAnnotationByName("Override").isPresent();
                if (!methodDeclaration.isAbstract() && !hasOverride) {
                    methodDeclaration.addAnnotation(new MarkerAnnotationExpr("Override"));
                }
                NodeList<Parameter> parameters = methodDeclaration.getParameters();
                List<String> curParams = parameters.stream().map(Parameter::getNameAsString).toList();
                for (int i = 0; i < insertIndexList.size(); i++) {
                    Type paramsType = paramsTypeList.get(i);
                    String paramName = RandomExpression.getName() + "_" + i;
                    while (curParams.contains(paramName)) {
                        paramName = RandomExpression.getName() + i;
                    }
                    params.add(new IMethod.AddTypeInfo(insertIndexList.get(i), paramsType, paramName));
                    parameters.add(insertIndexList.get(i), new Parameter(paramsType, paramName));
                }
                methodDeclaration.setParameters(parameters);
                //4.2.对函数添加与主逻辑无关的代码片段，用于扰乱代码逻辑
                if (!params.isEmpty()) {
                    methodDeclaration.getBody().ifPresent(blockStmt -> {
                        Map<String, String> paramInfos = params.stream().collect(Collectors.toMap(key -> key.paramName, value -> value.paramsType.asString()));
                        List<Statement> addStatements = RandomExpression.generateStatement(paramInfos);
                        NodeList<Statement> thisStatements = blockStmt.getStatements();
                        int addIndex = -1;
                        for (Statement statement : addStatements) {
                            int insertIndex = thisStatements.isEmpty() ? 0 : random.nextInt(thisStatements.size());
                            addIndex = Math.max(insertIndex, addIndex + 1);
                            thisStatements.add(addIndex, statement);
                        }
                        //调用工具类中的无效函数
                        int insertIndex = thisStatements.isEmpty() ? 0 : random.nextInt(thisStatements.size());
                        String returnName = RandomExpression.getName() + RandomExpression.getName();
                        String staticMethodCode = staticMethod.className + "." + staticMethod.md.getName() + "(";
                        List<String> thisParamNames = new ArrayList<>();
                        for (int num = 0; num < staticParams.size(); num++) {
                            thisParamNames.add(params.get(num).paramName);
                        }
                        staticMethodCode = staticMethodCode + String.join(", ", thisParamNames) + ");";
                        staticMethodCode = "Object " + returnName + " = " + staticMethodCode;
                        thisStatements.add(insertIndex, StaticJavaParser.parseStatement(staticMethodCode));
                        String returnPrint = RandomExpression.println + "(\"fun" + returnName + "\", " + returnName + ");";
                        thisStatements.add(insertIndex + 1, StaticJavaParser.parseStatement(returnPrint));
                        blockStmt.setStatements(thisStatements);
                    });
                }
                //4.3.对引用点添加参数信息
                List<MethodCallExpr> methodCallExprList = methodCallExprMap.get(ancestorMethod);
                if (null != methodCallExprList) {
                    for (MethodCallExpr mc : methodCallExprList) {
                        NodeList<Expression> arguments = mc.getArguments();
                        for (IMethod.AddTypeInfo aty : params) {
                            arguments.add(aty.insertIndex, RandomExpression.generateValue(aty.paramsType.asString()));
                        }
                        mc.setArguments(arguments);
                    }
                }
            }
        }
        //6.保存修改的文件
        ObfuscatorUtils.saveObfuscatedFiles(fileMap);
    }

    //是否启动这个功能
    @Override
    public boolean supports(ObfuscatorConfig config) {
        return ObjectUtil.isNotEmpty(config.getOverrideMethodParams());
    }


    /**
     * 遍历出哪些函数可以进行参数修改
     **/
    @Override
    public void visit(ClassOrInterfaceDeclaration cid, OverrideMethodArgumentParams config) {
        super.visit(cid, config);
        //保存所有的方法函数
        for (MethodDeclaration md : cid.getMethods()) {
            try {
                String methodSignature = md.resolve().getQualifiedSignature();
                methodDeclarationMap.put(methodSignature, md);
            } catch (Exception ignore) {

            }
        }
        //函数不是接口与内部类
        if (!cid.isInnerClass() && !cid.isInterface()) {
            //获取父类函数中的所有函数
            Map<String, List<String>> ams = new LinkedHashMap<>();
            try {
                for (ResolvedReferenceType ancestor : cid.resolve().getAncestors()) {
                    for (ResolvedMethodDeclaration x : ancestor.getAllMethods()) {
                        if (x.getQualifiedSignature().startsWith(config.getBasePackage())) {
                            if (!x.getName().equalsIgnoreCase(RandomExpression.getLoggerPrintName())) {
                                ams.computeIfAbsent(x.getSignature(), k -> new ArrayList<>()).add(x.getQualifiedSignature());
                            }
                        }
                    }
                }
                //获取当前类中的函数
                for (MethodDeclaration md : cid.getMethods()) {
                    for (String signature : ams.keySet()) {
                        if (signature.equals(md.resolve().getSignature())) {
                            if (!md.getNameAsString().equalsIgnoreCase(RandomExpression.getLoggerPrintName())) {
                                ams.get(signature).add(0, md.resolve().getQualifiedSignature());
                            }
                        }
                    }
                }
            } catch (Exception ignored) {
                if (print) {
                    System.err.println(ignored.getMessage());
                }
            }
            //聚合所有相同实现的类到同一个地方
            for (List<String> values : ams.values()) {
                if (values.size() <= 1) {
                    continue;
                }
                boolean has = false;
                for (Set<String> ancestorMethods : ancestorMethodList) {
                    for (String value : values) {
                        if (ancestorMethods.contains(value)) {
                            has = true;
                            ancestorMethods.addAll(values);
                        }
                    }
                }
                if (!has) {
                    ancestorMethodList.add(new HashSet<>(values));
                }
            }
        }
    }

    /**
     * 这里处理函数调用时，对参数的修改
     **/
    @Override
    public void visit(MethodCallExpr mc, OverrideMethodArgumentParams config) {
        super.visit(mc, config);
        mc.toMethodCallExpr().ifPresent(resolvedMethod -> {
            try {
                String methodSignature = resolvedMethod.resolve().getQualifiedSignature();
                methodCallExprMap.computeIfAbsent(methodSignature, k -> new ArrayList<>()).add(mc);
            } catch (Exception ignored) {
                if (print) {
                    System.err.println(ignored.getMessage());
                }
            }
        });
    }


    /**
     * 随机生成工具类
     **/
    public List<IMethod.StaticMethod> randomStaticMethods(String basePackage) {
        //当前新增的包
        List<IMethod.StaticMethod> newStaticMethods = new ArrayList<>();
        String packageName = "z" + RandomUtil.randomString(RandomUtil.BASE_CHAR, 5);
        //原本就是2
        for (int i = 0; i < IMethod.INVALID_PACKAGE_NUM; i++) {
            Set<Integer> indexList = new HashSet<>();
            String className = "UT" + RandomExpression.getName().toLowerCase();
            String thisPackage = basePackage + "." + packageName;
            for (int j = 0; j < Math.max(random.nextInt(10), 5); j++) {
                indexList.add(random.nextInt(tempStaticMethods.size()));
            }
            CompilationUnit cu = new CompilationUnit();
            cu.setPackageDeclaration(thisPackage);
            cu.addClass(className).setPublic(true).setFinal(true);
            cu.getClassByName(className).ifPresent(cid -> {
                for (Integer index : indexList) {
                    cid.addMember(tempStaticMethods.get(index));
                    newStaticMethods.add(new IMethod.StaticMethod(thisPackage, className, tempStaticMethods.get(index)));
                }
            });
            try {
                Path savePath = Paths.get(this.sourceRoot, thisPackage.replace(".", "/"));
                if (!savePath.toFile().exists()) {
                    if (!savePath.toFile().mkdirs()) {
                        throw new RuntimeException("mkdirs error");
                    }
                }
                Files.writeString(Paths.get(savePath.toString(), className + ".java"), cu.toString());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        return newStaticMethods;
    }
}
